
from isaaclab.utils import configclass
from isaaclab.sensors.frame_transformer.frame_transformer_cfg import FrameTransformerCfg
from isaaclab.sensors.frame_transformer.frame_transformer_cfg import OffsetCfg 
from isaaclab.assets import  AssetBaseCfg, RigidObjectCfg
from isaaclab.sim.schemas.schemas_cfg import RigidBodyPropertiesCfg 
from isaaclab.sim.spawners.from_files.from_files_cfg import GroundPlaneCfg, UsdFileCfg
from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR
from isaaclab_assets.robots.franka import FRANKA_PANDA_CFG

from arm_lab.tasks.manager_based.lift.lift_env_cfg import  LiftEnvCfg
from arm_lab.tasks.manager_based.lift import mdp
import math

@configclass
class FrankaEnvCfg(LiftEnvCfg):
    
    def __post_init__(self):
        # post init of parent
        super().__post_init__()

        # ------------------------------Sence------------------------------
        # switch robot 
        self.scene.robot = FRANKA_PANDA_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
        self.scene.ee_frame.prim_path = "{ENV_REGEX_NS}/Robot/panda_link0"
        self.scene.ee_frame.target_frames = [
                FrameTransformerCfg.FrameCfg(
                    prim_path="{ENV_REGEX_NS}/Robot/panda_hand",
                    name="end_effector",
                    offset=OffsetCfg(
                        pos=[0.0, 0.0, 0.1034],  # 夹爪末端到末端关节坐标系的偏置
                    ),
                ),
            ]
        # object 
        self.scene.object = RigidObjectCfg(
            prim_path="{ENV_REGEX_NS}/Object",
            init_state=RigidObjectCfg.InitialStateCfg(pos=[0.5, 0, 0.035], rot=[1, 0, 0, 0]),
            spawn=UsdFileCfg(
                usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Blocks/DexCube/dex_cube_instanceable.usd",
                scale=(0.8, 0.8, 0.8),
                rigid_props=RigidBodyPropertiesCfg(
                    solver_position_iteration_count=16,
                    solver_velocity_iteration_count=1,
                    max_angular_velocity=1000.0,
                    max_linear_velocity=1000.0,
                    max_depenetration_velocity=5.0,
                    disable_gravity=False,
                ),
            ),
        )
        # table 
        self.scene.table = AssetBaseCfg(
            prim_path="{ENV_REGEX_NS}/Table",
            init_state=AssetBaseCfg.InitialStateCfg(pos=[0.5, 0, 0], rot=[0.707, 0, 0, 0.707]),
            spawn=UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/SeattleLabTable/table_instanceable.usd"),
        )
        # plane
        self.scene.plane = AssetBaseCfg(
            prim_path="/World/GroundPlane",
            init_state=AssetBaseCfg.InitialStateCfg(pos=[0, 0, -1.05]),
            spawn=GroundPlaneCfg(),
        )

        # ------------------------------Commands----------------------------------- 
        self.commands.object_pose.body_name = "panda_hand"
        self.commands.object_pose.ranges = mdp.UniformPoseCommandCfg.Ranges(
            pos_x=(0.4, 0.6), pos_y=(-0.25, 0.25), pos_z=(0.25, 0.5), roll=(0.0, 0.0), pitch=(0.0, 0.0), yaw=(0.0, 0.0)
        )

        # ------------------------------Actions-----------------------------------
        self.actions.arm_action.joint_names = ["panda_joint.*"]
        self.actions.gripper_action.joint_names = ["panda_finger.*"]
        self.actions.gripper_action.open_command_expr = {"panda_finger_.*": 0.04}
        self.actions.gripper_action.close_command_expr = {"panda_finger_.*": 0.0}

        # ------------------------------Observations------------------------------

        # ------------------------------Events------------------------------------
        self.events.randomize_apply_external_force_torque = None
        self.events.randomize_rigid_body_mass = None
        self.events.randomize_rigid_body_inertia = None

        # ------------------------------Rewards------------------------------
        # tracking reward
        self.rewards.reaching_object.weight = 2.0
        self.rewards.lifting_object.weight = 15.0
        self.rewards.object_goal_tracking.weight = 16.0
        self.rewards.object_goal_tracking_fine_grained.weight = 5.0

        # action penalty
        self.rewards.action_rate.weight = -1e-1
        self.rewards.joint_vel.weight = -1e-1
        self.rewards.ee_vel.weight = - 1.0

        # ------------------------------Terminations------------------------------
        
        # set rewards which weight 0.0 to None
        # If the weight of rewards is 0, set rewards to None
        if self.__class__.__name__ == "FrankaEnvCfg":
            self.disable_zero_weight_rewards()
        
